package main

import (
	"context"
	"fmt"
	"log"
	"time"

	"github.com/louloulin/dataflare/pkg/state/query"
	"github.com/louloulin/dataflare/pkg/state/store"
	"github.com/louloulin/dataflare/pkg/state/store/memory"
	"github.com/louloulin/dataflare/pkg/state/transaction"
)

func main() {
	// Create a state store
	stateStore := memory.NewMemoryStateStore()
	ctx := context.Background()
	err := stateStore.Init(ctx, nil)
	if err != nil {
		log.Fatalf("Failed to initialize state store: %v", err)
	}

	// Create a KV state
	log.Println("Creating KV state...")
	kvConfig := &store.StateConfig{
		Name:       "users",
		Type:       store.StateTypeKV,
		Scope:      store.StateScopeLocal,
		TTL:        time.Hour,
		Persistent: false,
	}
	kvState, err := stateStore.CreateState(ctx, kvConfig)
	if err != nil {
		log.Fatalf("Failed to create KV state: %v", err)
	}

	// Create a counter state
	log.Println("Creating counter state...")
	counterConfig := &store.StateConfig{
		Name:       "user-count",
		Type:       store.StateTypeCounter,
		Scope:      store.StateScopeLocal,
		TTL:        time.Hour,
		Persistent: false,
	}
	counterState, err := stateStore.CreateState(ctx, counterConfig)
	if err != nil {
		log.Fatalf("Failed to create counter state: %v", err)
	}

	// Create a transaction manager
	log.Println("Creating transaction manager...")
	txManager := transaction.NewTransactionManager(stateStore)

	// Begin a transaction
	log.Println("Beginning transaction...")
	tx, err := txManager.Begin(ctx, "tx1", nil)
	if err != nil {
		log.Fatalf("Failed to begin transaction: %v", err)
	}

	// Add some users
	log.Println("Adding users...")
	err = tx.Set(ctx, "users", "user1", []byte(`{"id":"user1","name":"Alice","age":30}`))
	if err != nil {
		log.Fatalf("Failed to set user: %v", err)
	}
	err = tx.Set(ctx, "users", "user2", []byte(`{"id":"user2","name":"Bob","age":25}`))
	if err != nil {
		log.Fatalf("Failed to set user: %v", err)
	}
	err = tx.Set(ctx, "users", "user3", []byte(`{"id":"user3","name":"Charlie","age":35}`))
	if err != nil {
		log.Fatalf("Failed to set user: %v", err)
	}

	// Commit the transaction
	log.Println("Committing transaction...")
	err = txManager.Commit(ctx, "tx1")
	if err != nil {
		log.Fatalf("Failed to commit transaction: %v", err)
	}

	// Update the user count
	log.Println("Updating user count...")
	counterState.(store.CounterState).Increment(ctx, "total", 3)

	// Create a query engine
	log.Println("Creating query engine...")
	queryEngine := query.NewQueryEngine(stateStore)

	// Get a user
	log.Println("Getting a user...")
	getQuery := &query.Query{
		Type:  query.QueryTypeGet,
		State: "users",
		Key:   "user1",
	}
	getResult, err := queryEngine.Execute(ctx, getQuery)
	if err != nil {
		log.Fatalf("Failed to execute get query: %v", err)
	}
	log.Printf("User: %s", string(getResult.Records[0]["value"].([]byte)))

	// List all users
	log.Println("Listing all users...")
	listQuery := &query.Query{
		Type:  query.QueryTypeList,
		State: "users",
	}
	listResult, err := queryEngine.Execute(ctx, listQuery)
	if err != nil {
		log.Fatalf("Failed to execute list query: %v", err)
	}
	log.Printf("Found %d users", listResult.Count)
	for _, record := range listResult.Records {
		log.Printf("User: %s", string(record["value"].([]byte)))
	}

	// Search for users
	log.Println("Searching for users...")
	searchQuery := &query.Query{
		Type:  query.QueryTypeSearch,
		State: "users",
		Condition: &query.Condition{
			Field:              "key",
			ComparisonOperator: query.ComparisonOperatorLike,
			Value:              "user%",
		},
		Sort: &query.Sort{
			Field: "key",
			Order: query.SortOrderAscending,
		},
	}
	searchResult, err := queryEngine.Execute(ctx, searchQuery)
	if err != nil {
		log.Fatalf("Failed to execute search query: %v", err)
	}
	log.Printf("Found %d users", searchResult.Count)
	for _, record := range searchResult.Records {
		log.Printf("User: %s", string(record["value"].([]byte)))
	}

	// Count users
	log.Println("Counting users...")
	countQuery := &query.Query{
		Type:  query.QueryTypeAggregate,
		State: "users",
		Aggregate: &query.Aggregate{
			Function: query.AggregateFunctionCount,
			Field:    "key",
		},
	}
	countResult, err := queryEngine.Execute(ctx, countQuery)
	if err != nil {
		log.Fatalf("Failed to execute count query: %v", err)
	}
	log.Printf("User count: %v", countResult.AggregateValue)

	// Get the user count from the counter state
	log.Println("Getting user count from counter state...")
	count, err := counterState.(store.CounterState).Get(ctx, "total")
	if err != nil {
		log.Fatalf("Failed to get user count: %v", err)
	}
	log.Printf("User count from counter: %d", count)

	// Begin another transaction
	log.Println("Beginning another transaction...")
	tx2, err := txManager.Begin(ctx, "tx2", nil)
	if err != nil {
		log.Fatalf("Failed to begin transaction: %v", err)
	}

	// Update a user
	log.Println("Updating a user...")
	err = tx2.Set(ctx, "users", "user2", []byte(`{"id":"user2","name":"Bob","age":26}`))
	if err != nil {
		log.Fatalf("Failed to set user: %v", err)
	}

	// Delete a user
	log.Println("Deleting a user...")
	err = tx2.Delete(ctx, "users", "user3")
	if err != nil {
		log.Fatalf("Failed to delete user: %v", err)
	}

	// Commit the transaction
	log.Println("Committing transaction...")
	err = txManager.Commit(ctx, "tx2")
	if err != nil {
		log.Fatalf("Failed to commit transaction: %v", err)
	}

	// Update the user count
	log.Println("Updating user count...")
	counterState.(store.CounterState).Decrement(ctx, "total", 1)

	// List all users again
	log.Println("Listing all users again...")
	listResult, err = queryEngine.Execute(ctx, listQuery)
	if err != nil {
		log.Fatalf("Failed to execute list query: %v", err)
	}
	log.Printf("Found %d users", listResult.Count)
	for _, record := range listResult.Records {
		log.Printf("User: %s", string(record["value"].([]byte)))
	}

	// Get the user count from the counter state again
	log.Println("Getting user count from counter state again...")
	count, err = counterState.(store.CounterState).Get(ctx, "total")
	if err != nil {
		log.Fatalf("Failed to get user count: %v", err)
	}
	log.Printf("User count from counter: %d", count)

	// Clean up transactions
	log.Println("Cleaning up transactions...")
	err = txManager.CleanUp(ctx, 0)
	if err != nil {
		log.Fatalf("Failed to clean up transactions: %v", err)
	}

	// Close the state store
	log.Println("Closing state store...")
	err = stateStore.Close(ctx)
	if err != nil {
		log.Fatalf("Failed to close state store: %v", err)
	}

	fmt.Println("State management example completed successfully!")
}
